import logging
import os
from math import sqrt
from numbers import Real

import numpy as np
from PIL import Image
from torchvision import transforms

from simpletuner.helpers.training.state_tracker import StateTracker

logger = logging.getLogger("MultiaspectImage")
logger.setLevel(os.environ.get("SIMPLETUNER_IMAGE_PREP_LOG_LEVEL", "INFO"))


class MultiaspectImage:
    @staticmethod
    def _coerce_positive_int(value, default: int = 1) -> int:
        """Return a positive integer from value or fallback to default."""
        if isinstance(value, Real):
            result = int(value)
        else:
            try:
                result = int(value)
            except (TypeError, ValueError):
                return default
        return result if result > 0 else default

    @staticmethod
    def _get_alignment(default: int = 1) -> int:
        args = StateTracker.get_args()
        align = getattr(args, "aspect_bucket_alignment", None) if args is not None else None
        return MultiaspectImage._coerce_positive_int(align, default)

    @staticmethod
    def _get_rounding(default: int) -> int:
        args = StateTracker.get_args()
        rounding = getattr(args, "aspect_bucket_rounding", None) if args is not None else None
        if isinstance(rounding, Real):
            rounding = int(rounding)
        else:
            try:
                rounding = int(rounding)
            except (TypeError, ValueError):
                rounding = None
        if rounding is None or rounding < 0:
            return default
        return rounding

    @staticmethod
    def limit_canvas_size(width: int, height: int, max_size: int) -> dict:
        """
        Limit the canvas size to a maximum value.
        If the canvas size exceeds the maximum, it will be reduced proportionally.
        It's important that the canvas size returned by this function is used for cropping,
          since the aspect ratio of the canvas size is not guaranteed to match the
          aspect ratio of the original image.

        Args:
            width (int): The width of the image.
            height (int): The height of the image.
            max_size (int): The maximum allowed canvas size.

        Returns:
            dict: A dictionary containing the adjusted width, height, and canvas size.
        """
        if width * height <= max_size:
            # If the canvas size is already within limits, return the original dimensions.
            return {"width": width, "height": height, "canvas_size": width * height}

        align = MultiaspectImage._get_alignment()
        dims = [("width", width), ("height", height)]

        # Sort by size descending
        dims.sort(key=lambda x: x[1], reverse=True)
        # Subtract alignment from the larger dimension
        if dims[0][0] == "width":
            width -= align
        else:
            height -= align

        # If still too large, subtract from the other dimension.
        # Note: we're sorting by the original width and height values here to ensure we adjust the originally-smaller-dimension.
        if width * height > max_size:
            if dims[1][0] == "width":
                width -= align
            else:
                height -= align

        return {"width": width, "height": height, "canvas_size": width * height}

    @staticmethod
    def _round_to_nearest_multiple(value, override_value: int = None):
        """Round a value to the nearest multiple."""
        multiple = override_value
        if multiple is None:
            multiple = MultiaspectImage._get_alignment()
        else:
            multiple = MultiaspectImage._coerce_positive_int(multiple, default=1)
        rounded = round(value / multiple) * multiple
        return max(rounded, multiple)  # Ensure it's at least the value of 'multiple'

    @staticmethod
    def is_image_too_large(image_size: tuple, resolution: float, resolution_type: str):
        """
        Determine if an image is too large to be processed.

        Args:
            image (PIL.Image): The image to check.
            resolution (float): The maximum resolution to allow.
            resolution_type (str): What form of resolution to check, choices: "pixel", "area".

        Returns:
            bool: True if the image is too large, False otherwise.
        """
        if resolution_type == "pixel":
            return image_size[0] > resolution or image_size[1] > resolution
        elif resolution_type == "area":
            image_area = image_size[0] * image_size[1]
            target_area = resolution * 1e6  # Convert megapixels to pixels
            logger.debug(
                f"Image is too large? {image_area > target_area} (image area: {image_area}, target area: {target_area})"
            )
            return image_area > target_area
        else:
            raise ValueError(f"Unknown resolution type: {resolution_type}")

    @staticmethod
    def calculate_new_size_by_pixel_edge(aspect_ratio: float, resolution: int, original_size: tuple):
        if type(aspect_ratio) != float:
            raise ValueError(f"Aspect ratio must be a float, not {type(aspect_ratio)}")
        if type(resolution) != int and (type(resolution) != float or int(resolution) != resolution):
            raise ValueError(f"Resolution must be an int, not {type(resolution)}")

        W_original, H_original = original_size
        # Reject invalid images with extreme dimensions
        if W_original <= 0 or H_original <= 0:
            raise ValueError(f"Invalid image dimensions: {original_size}")

        # Use raw (unrounded) aspect ratio for calculations to avoid division by zero
        original_aspect_ratio_raw = W_original / H_original

        # Calculate target size based on TARGET aspect for final crop
        if W_original < H_original:  # Portrait or square orientation
            W_target = resolution
            H_target = int(W_target / aspect_ratio)
        else:  # Landscape orientation
            H_target = resolution
            W_target = int(H_target * aspect_ratio)

        # Round to nearest multiple
        W_adjusted = MultiaspectImage._round_to_nearest_multiple(W_target)
        H_adjusted = MultiaspectImage._round_to_nearest_multiple(H_target)

        # Calculate intermediary size that maintains ORIGINAL aspect ratio
        # and is large enough to crop to target size
        if original_aspect_ratio_raw >= 1.0:  # Landscape or square original
            # Make height match the target's height requirement, width follows original aspect
            H_initial = max(H_adjusted, H_target)
            W_initial = int(H_initial * original_aspect_ratio_raw)
            # Ensure width is also large enough
            if W_initial < W_adjusted:
                W_initial = W_adjusted
                H_initial = int(W_initial / original_aspect_ratio_raw)
        else:  # Portrait original
            # Make width match the target's width requirement, height follows original aspect
            W_initial = max(W_adjusted, W_target)
            H_initial = int(W_initial / original_aspect_ratio_raw)
            # Ensure height is also large enough
            if H_initial < H_adjusted:
                H_initial = H_adjusted
                W_initial = int(H_initial * original_aspect_ratio_raw)

        adjusted_aspect_ratio = MultiaspectImage.calculate_image_aspect_ratio((W_adjusted, H_adjusted))

        return (W_adjusted, H_adjusted), (W_initial, H_initial), adjusted_aspect_ratio

    @staticmethod
    def calculate_new_size_by_pixel_area(aspect_ratio: float, megapixels: float, original_size: tuple):
        if type(aspect_ratio) not in [float, np.float64]:
            raise ValueError(f"Aspect ratio must be a float, not {type(aspect_ratio)}")
        target_pixel_area = megapixels * 1e6  # Convert megapixels to pixel area, eg. 1.0 mp = 1000000 pixels
        target_pixel_edge = MultiaspectImage._round_to_nearest_multiple(int(sqrt(target_pixel_area)))
        logger.debug(
            f"Converted {megapixels} megapixels to {target_pixel_area} pixels with a square edge of {target_pixel_edge}."
        )

        W_initial, H_initial = original_size
        # Reject invalid images with extreme dimensions
        if W_initial <= 0 or H_initial <= 0:
            raise ValueError(f"Invalid image dimensions: {original_size}")

        if aspect_ratio == 1.0 and W_initial == H_initial:
            # If the aspect ratio is 1.0 and original is square, resize straight to target.
            logger.debug(
                f"Returning the square edge {target_pixel_edge}x{target_pixel_edge} as both target and intermediary for square input."
            )
            return (
                (target_pixel_edge, target_pixel_edge),
                (target_pixel_edge, target_pixel_edge),
                aspect_ratio,
            )

        # Calculate the target size. This is what will be cropped-to.
        W_target = MultiaspectImage._round_to_nearest_multiple(target_pixel_edge * sqrt(aspect_ratio))
        H_target = MultiaspectImage._round_to_nearest_multiple(target_pixel_edge / sqrt(aspect_ratio))
        calculated_resulting_megapixels = (W_target * H_target) / 1e6
        target_aspect_ratio = MultiaspectImage.calculate_image_aspect_ratio((W_target, H_target))

        if not np.isclose(calculated_resulting_megapixels, megapixels, rtol=1e-1):
            logger.debug(
                f"-!- This image will not have the correct target megapixel size: {calculated_resulting_megapixels}"
            )

        # Calculate the intermediary size. This will maintain ORIGINAL aspect ratio and be resized-to.
        # Use raw (unrounded) aspect ratio for calculations to avoid division by zero
        original_aspect_ratio_raw = W_initial / H_initial
        if W_target < H_target:  # Portrait or square orientation
            W_intermediary = W_target
            H_intermediary = int(W_intermediary / original_aspect_ratio_raw)
        else:  # Landscape orientation
            H_intermediary = H_target
            W_intermediary = int(H_intermediary * original_aspect_ratio_raw)

        # retrieve the static mapping.
        adjusted_aspect_ratio = MultiaspectImage.calculate_image_aspect_ratio((W_target, H_target))
        previously_stored_resolution = StateTracker.get_resolution_by_aspect(
            dataloader_resolution=megapixels, aspect=adjusted_aspect_ratio
        )

        if previously_stored_resolution:
            logger.debug(
                f"Using cached aspect-resolution map value for {adjusted_aspect_ratio}: {previously_stored_resolution}"
            )
            W_target, H_target = previously_stored_resolution
        target_resolution = (W_target, H_target)

        # The intermediary size might be smaller than the target. This is bad.
        # If it happens, the cropped image will be cropped past the boundaries of the intermediary size.
        if W_target > W_intermediary or H_target > H_intermediary:
            _W_intermediary, _H_intermediary = W_intermediary, H_intermediary
            if W_target > W_intermediary:
                W_diff = W_target - W_intermediary
                H_diff = int(W_diff / original_aspect_ratio_raw)
            else:
                H_diff = H_target - H_intermediary
                W_diff = int(H_diff * original_aspect_ratio_raw)
            H_intermediary += H_diff
            W_intermediary += W_diff
            logger.debug(
                f"Intermediary size {_W_intermediary}x{_H_intermediary} would be smaller than {W_target}x{H_target} with a difference in size of {W_diff}x{H_diff}."
                f" The size will be adjusted to maintain the original aspect ratio: {W_intermediary}x{H_intermediary}."
            )
            calculated_resulting_megapixels = (W_intermediary * H_intermediary) / 1e6

        intermediary_resolution = (W_intermediary, H_intermediary)

        logger.debug(
            f"Using target size of {megapixels} megapixels:"
            f"\n-> initial size is {W_initial}x{H_initial}, original aspect ratio {aspect_ratio}."
            f"\n-> intermediary size is {W_intermediary}x{H_intermediary}, with aspect ratio {adjusted_aspect_ratio}."
            f"\n-> cropped size is {W_target}x{H_target}, with aspect ratio {target_aspect_ratio}."
            f"\n-> cropped sample will be {calculated_resulting_megapixels} megapixels"
        )
        # Attempt to retrieve previously stored resolution by adjusted aspect ratio
        if not previously_stored_resolution:
            logger.debug(
                f"No cached resolution found for aspect ratio {adjusted_aspect_ratio}. Storing {target_resolution}."
            )
            StateTracker.set_resolution_by_aspect(
                dataloader_resolution=megapixels,
                aspect=adjusted_aspect_ratio,
                resolution=target_resolution,
            )

        return (target_resolution, intermediary_resolution, adjusted_aspect_ratio)

    @staticmethod
    def adjust_resolution_to_bucket_interval(initial_resolution: tuple, target_resolution: tuple):
        W_initial, H_initial = initial_resolution
        W_adjusted, H_adjusted = target_resolution
        # If W_initial or H_initial are < W_adjusted or H_adjusted, add the greater of the two differences to both values.
        W_diff = W_adjusted - W_initial
        H_diff = H_adjusted - H_initial
        if W_diff > 0 and (W_diff > H_diff or W_diff == H_diff):
            logger.debug(
                f"Intermediary size {W_initial}x{H_initial} would be smaller than {W_adjusted}x{H_adjusted} with a difference in size of {W_diff}x{H_diff}. Adjusting both sides by {max(W_diff, H_diff)} pixels."
            )
            H_initial += W_diff
            W_initial += W_diff
        elif H_diff > 0 and H_diff > W_diff:
            logger.debug(
                f"Intermediary size {W_initial}x{H_initial} would be smaller than {W_adjusted}x{H_adjusted} with a difference in size of {W_diff}x{H_diff}. Adjusting both sides by {max(W_diff, H_diff)} pixels."
            )
            W_initial += H_diff
            H_initial += H_diff

        return W_initial, H_initial

    @staticmethod
    def calculate_image_aspect_ratio(image, rounding: int = 2):
        """
        Calculate the aspect ratio of an image and round it to a specified precision.

        Args:
            image (PIL.Image): The image to calculate the aspect ratio for.

        Returns:
            float: The rounded aspect ratio of the image.
        """
        to_round = MultiaspectImage._get_rounding(rounding)
        if isinstance(image, Image.Image):
            # An actual image was passed in.
            width, height = image.size
        elif isinstance(image, tuple) or isinstance(image, list):
            # An image.size or a similar (W, H) tuple was provided.
            width, height = image
        elif isinstance(image, float):
            # An externally-calculated aspect ratio was given to round.
            return round(image, to_round)
        elif isinstance(image, np.ndarray):
            # A video was passed in as a numpy array.
            width, height = image.shape[2], image.shape[1]
        else:
            raise ValueError(f"Unexpected type {image}")
            width, height = image.size
        aspect_ratio = round(width / height, to_round)
        return aspect_ratio

    @staticmethod
    def numpy_list_to_pil(numpy_list):
        if isinstance(numpy_list, list) and isinstance(numpy_list[0], np.ndarray):
            numpy_list = [Image.fromarray(np.uint8(image)).convert("RGB") for image in numpy_list]
        return numpy_list


resize_helpers = {
    "pixel": MultiaspectImage.calculate_new_size_by_pixel_edge,
    "area": MultiaspectImage.calculate_new_size_by_pixel_area,
}
